# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import json
import time
from decimal import Decimal
from typing import List
from tortoise.queryset import Q
from tortoise.transactions import in_transaction
from hypertext import PagingResult
from exception import AppException
from common.utils.urls import UrlUtil
from common.utils.tools import ToolsUtil
from common.enums.agent import AgentEnum
from common.enums.wallet import WalletEnum
from common.enums.models import ModelsEnum
from common.models.users import UserModel
from common.models.users import UserWalletModel
from common.models.dev import DevModelsModel
from common.models.agents import AgentKnowModel
from common.models.agents import AgentKnowArchiveModel
from common.postgres.public import PgKnowledgeModel
from common.postgres.public import PgAttachmentModel
from common.cache.pool_cache import KeyPoolCache
from common.chain.chain_server import ChatUtils
from common.chain.vector_server import VectorService
from common.chain.chain_server import AIChannelFactory, FlowsSchema
from common.service.retrieval_logic import RecallKnowLogic, RecallUtil
from apps.api.schemas.agent.know import dataset_schema as schema
from apps.api.service.agent.know.teamed_service import KnowTeamedService


class KnowDatasetService:
    """ 文档数据服务类 """

    @classmethod
    async def lists(cls, user_id: int, params: schema.KnowDatasetSearchIn) -> PagingResult[schema.KnowDatasetListVo]:
        """
        文档数据列表

        Args:
            user_id (int): 用户ID
            params (schema.KnowDatasetSearchIn): 数据搜索参数

        Returns:
            PagingResult[schema.KnowDatasetListVo]

        Author:
            zero
        """
        # 查找文档
        archive = await AgentKnowArchiveModel.filter(code=params.archive_id).first()
        if not archive:
            raise AppException("文档不存在哦")
        if archive.is_delete:
            raise AppException("文档已被删除")

        # 验知识库
        know = await AgentKnowModel.filter(id=archive.know_id).first()
        if not know:
            raise AppException("知识库已不存在")
        if know.is_delete:
            raise AppException("知识库已被删除")
        if know.is_disable:
            raise AppException("知识库已被禁用")

        # 验证权限
        if user_id != know.user_id:
            await KnowTeamedService.has_permission(know.code, user_id, AgentEnum.POWER_VIEW)

        # 搜索条件
        where = [Q(archive_id=archive.id), Q(is_delete=0)]
        if params.keyword:
            where.append(Q(question__icontains=params.keyword))

        # 查询数据
        _model = PgKnowledgeModel.filter(*where).order_by("-create_time", "chunk_index")
        _pager = await PgKnowledgeModel.paginate(
            model=_model,
            page_no=params.page_no,
            page_size=params.page_size
        )

        # 处理数据
        _lists = []
        for item in _pager.lists:
            _lists.append(schema.KnowDatasetListVo(
                id=str(item.get("uuid")),
                question=item.get("question") or "",
                answer=item.get("answer") or "",
                dimension=item.get("dimension"),
                character=len(item.get("question") or "") + len(item.get("answer") or ""),
                status=item["status"],
                is_disable=item["is_disable"],
                create_time=item["create_time"],
                update_time=item["create_time"],
            ))

        _pager.lists = _lists
        return _pager

    @classmethod
    async def detail(cls, user_id: int, uuid_: str) -> schema.KnowDatasetDetailVo:
        """
        数据详情信息

        Args:
            user_id (int): 用户ID
            uuid_ (str): 数据的UUID

        Returns:
            schema.KnowDatasetDetailVo

        Author:
            zero
        """
        # 查询数据
        detail = await PgKnowledgeModel.filter(uuid=uuid_).get()

        # 验知识库
        know = await AgentKnowModel.filter(id=detail.know_id).first()
        if not know:
            raise AppException("知识库已不存在")
        if know.is_delete:
            raise AppException("知识库已被删除")
        if know.is_disable:
            raise AppException("知识库已被禁用")

        # 验证权限
        if user_id != know.user_id:
            await KnowTeamedService.has_permission(know.code, user_id, AgentEnum.POWER_VIEW)

        # 获取附件
        annex = []
        attachment = await (PgAttachmentModel
                            .filter(dataset_id=detail.uuid)
                            .filter(is_delete=0)
                            .order_by("sort").all())
        for att in attachment:
            annex.append(schema.KnowDatasetAnnex(
                id=str(att.uuid),
                type=att.type,
                name=att.name,
                path=await UrlUtil.to_absolute_url(att.path),
                size=att.size,
                ext=att.ext,
                sort=att.sort,
                question=att.question
            ))

        # 返回结果
        return schema.KnowDatasetDetailVo(
            id=str(uuid_),
            question=detail.question,
            answer=detail.answer,
            annex=annex
        )

    @classmethod
    async def insert(cls, user_id: int, post: schema.KnowDatasetInsertIn):
        """
        文档数据插入

        Args:
            user_id (int): 用户ID
            post (schema.KnowDatasetInsertIn): 插入参数

        Author:
            zero
        """
        # 查找文档
        archive = await AgentKnowArchiveModel.filter(code=post.fid).first()
        if not archive:
            raise AppException("文档不存在哦")
        if archive.is_delete:
            raise AppException("文档已被删除")

        # 验知识库
        know = await AgentKnowModel.filter(id=int(archive.know_id)).first()
        if not know:
            raise AppException("知识库已不存在")
        if know.is_delete:
            raise AppException("知识库已被删除")
        if know.is_disable:
            raise AppException("知识库已被禁用")

        # 验证权限
        if user_id != know.user_id:
            await KnowTeamedService.has_permission(know.code, user_id, AgentEnum.POWER_EDIT)

        # 累计消耗
        total_use_tokens: int = 0
        total_use_points: int = 0
        additional = []

        # 问题向量
        start_time = time.time()
        emb: dict = await cls._post_embedding(know, [post.question])
        ask_emb: str = "[" + ",".join(str(item) for item in emb["embeddings"][0]) + "]"
        total_use_tokens += emb["use_tokens"]
        total_use_points += emb["use_points"]
        additional.append(FlowsSchema(
            name="索引生成",
            scene="vector",
            alias=emb["alias"],
            model=emb["model"],
            task_time=f"{(time.time() - start_time):.2f}",
            use_points=emb["use_points"],
            total_tokens=emb["use_tokens"],
            prompt_tokens=emb["use_tokens"],
            completion_tokens=0
        ))

        # 附件向量
        start_time = time.time()
        texts = [(item.question or item.name).strip() for item in post.annex]
        attrs = await cls._post_embedding(know, texts)
        if texts:
            total_use_tokens += attrs["use_tokens"]
            total_use_points += attrs["use_points"]
            additional.append(FlowsSchema(
                    name="附件索引",
                    scene="vector",
                    alias=attrs["alias"],
                    model=attrs["model"],
                    task_time=f"{(time.time() - start_time):.2f}",
                    use_points=attrs["use_points"],
                    total_tokens=attrs["use_tokens"],
                    prompt_tokens=attrs["use_tokens"],
                    completion_tokens=0
                ))

        # 更新数据
        async with in_transaction("pgsql"):
            # 插入数据
            text: str = post.question.strip() + post.answer.strip()
            signed: str = ToolsUtil.make_md5_str(text)
            result = await PgKnowledgeModel.create(
                know_id=know.id,
                archive_id=archive.id,
                user_id=user_id,
                last_uid=user_id,
                question=post.question,
                answer=post.answer,
                signed=signed,
                embedding=ask_emb,
                use_tokens=0,
                status=AgentEnum.EMB_YES,
                create_time=int(time.time()),
                update_time=int(time.time())
            )

            # 更新分词
            await PgKnowledgeModel.update_ts_vectory([str(result.uuid)])

            # 处理附件
            index: int = 0
            annex = []
            for item in post.annex:
                texts = (item.question or item.name).strip()
                signed = ToolsUtil.make_md5_str(texts)
                annex.append(PgAttachmentModel(
                    scene="know",
                    dataset_id=str(result.uuid),
                    archive_id=archive.id,
                    know_id=know.id,
                    type=item.type,
                    name=item.name,
                    path=item.path,
                    ext=item.ext,
                    size=item.size,
                    sort=item.sort,
                    signed=signed,
                    question=texts,
                    status=AgentEnum.EMB_YES,
                    vector_alias=attrs["alias"],
                    vector_model=attrs["model"],
                    use_tokens=attrs["use_tokens"],
                    dimension=len(attrs["embeddings"][index]),
                    embedding=attrs["embeddings"][index],
                    create_time=int(time.time()),
                    update_time=int(time.time())
                ))
                index += 1

            if annex:
                await PgAttachmentModel.bulk_create(annex)
                await PgAttachmentModel.update_ts_vectory(str(result.uuid))

        # 扣费逻辑
        async with in_transaction("mysql"):
            # 用户扣费
            if total_use_points > 0:
                user = await UserModel.filter(id=know.id).first()
                surplus_balance: int = max(user.points - total_use_points, 0)
                await UserModel.filter(id=user.id).update(points=surplus_balance)

            # 写入日志
            if attrs.get("price"):
                flows_usage = [flow.model_dump() for flow in additional]
                await UserWalletModel.dec(
                    scene=WalletEnum.POINTS,
                    user_id=user.id,
                    change_amount=total_use_points,
                    change_type=WalletEnum.POINTS_DEC_KB_EMB,
                    remarks=WalletEnum.get_change_type_msg(WalletEnum.POINTS_DEC_KB_EMB),
                    project=archive.name,
                    source_sn=archive.code,
                    source_id=archive.id,
                    additional=json.dumps(flows_usage, ensure_ascii=False)
                )

    @classmethod
    async def update(cls, user_id: int, post: schema.KnowDatasetUpdateIn):
        """
        文档数据编辑

        Args:
            user_id (int): 用户ID
            post (schema.KnowDatasetUpdateIn): 插入参数

        Author:
            zero
        """
        # 查找数据
        dataset = await PgKnowledgeModel.filter(uuid=post.uuid).first()
        if not dataset:
            raise AppException("数据不存在")
        if dataset.is_delete:
            raise AppException("数据已被删")

        # 查找文档
        archive = await AgentKnowArchiveModel.filter(id=dataset.archive_id).first()
        if not archive:
            raise AppException("文档不存在")
        if archive.is_delete:
            raise AppException("文档已被删")

        # 验知识库
        know = await AgentKnowModel.filter(id=archive.know_id).first()
        if not know:
            raise AppException(str("知识库已不存在"))
        if know.is_disable:
            raise AppException(str("知识库已被删除"))
        if know.is_disable:
            raise AppException("知识库已被禁用")

        # 验证权限
        if user_id != know.user_id:
            await KnowTeamedService.has_permission(know.code, user_id, AgentEnum.POWER_EDIT)

        # 原有附件
        atts_old = await PgAttachmentModel.filter(dataset_id=post.uuid, is_delete=0).all()
        atts_arr = {str(item.uuid): item for item in atts_old}

        # 删除附件
        annex_old_ids = set([str(item.uuid) for item in atts_old])
        annex_new_ids = set([str(item.id) for item in post.annex if item.id])
        annex_delete = list(annex_old_ids - annex_new_ids)

        # 处理附件
        annex_create = []
        annex_update = []
        for item in post.annex:
            texts = (item.question or item.name).strip()
            signed = ToolsUtil.make_md5_str(texts)
            clz = atts_arr.get(item.id)
            if not clz:
                annex_create.append(PgAttachmentModel(
                    scene=str("know"),
                    dataset_id=post.uuid,
                    archive_id=archive.id,
                    last_uid=user_id,
                    know_id=know.id,
                    type=item.type,
                    name=item.name,
                    path=UrlUtil.to_relative_url(item.path),
                    ext=item.ext,
                    size=item.size,
                    sort=item.sort,
                    signed=signed,
                    question=texts,
                    create_time=int(time.time()),
                    update_time=int(time.time())
                ))
            else:
                if clz.signed != signed:
                    clz.ext = item.ext
                    clz.type = item.type
                    clz.size = item.size
                    clz.sort = item.sort
                    clz.signed = signed
                    clz.question = texts
                    clz.path = UrlUtil.to_relative_url(item.path)
                    clz.update_time = int(time.time())
                    annex_update.append(clz)

        # 累计消耗
        total_use_tokens: int = 0
        total_use_points: int = 0
        additional = []

        # 文本索引
        ask_emb: str = ""
        text: str = post.question.strip() + post.answer.strip()
        signed: str = ToolsUtil.make_md5_str(text)
        start_time = time.time()
        if dataset.signed != signed:
            emb: dict = await cls._post_embedding(know, [str(post.question)])
            ask_emb: str = "[" + ",".join(str(item) for item in emb["embeddings"][0]) + "]"
            total_use_tokens += emb["use_tokens"]
            total_use_points += emb["use_points"]
            additional.append(FlowsSchema(
                name=str("索引生成"),
                scene="vector",
                alias=emb["alias"],
                model=emb["model"],
                task_time=f"{(time.time() - start_time):.2f}",
                use_points=emb["use_points"],
                total_tokens=emb["use_tokens"],
                prompt_tokens=emb["use_tokens"],
                completion_tokens=0
            ))

        # 附件索引
        start_time = time.time()
        attr_create = await cls._post_embedding(know, [item.question for item in annex_create])
        attr_update = await cls._post_embedding(know, [item.question for item in annex_update])
        if attr_create["embeddings"] or attr_update["embeddings"]:
            use_tokens = attr_create["use_tokens"] + attr_update["use_tokens"]
            use_points = attr_create["use_points"] + attr_update["use_points"]
            total_use_tokens += use_tokens
            total_use_points += use_points
            additional.append(FlowsSchema(
                name="附件索引",
                scene="vector",
                alias=attr_create["alias"],
                model=attr_create["model"],
                task_time=f"{(time.time() - start_time):.2f}",
                use_points=use_points,
                total_tokens=use_tokens,
                prompt_tokens=use_tokens,
                completion_tokens=0
            ))

        async with in_transaction("pgsql"):
            # 更新数据
            if ask_emb:
                await PgKnowledgeModel.filter(uuid=post.uuid).update(
                    signed=signed,
                    question=post.question,
                    answer=post.answer,
                    update_time=int(time.time())
                )

            # 附件处理
            if annex_create:
                index: int = 0
                for item in annex_create:
                    item.vector_alias = attr_create["alias"]
                    item.vector_model = attr_create["model"]
                    item.dimension = len(attr_create["embeddings"][index])
                    item.embedding = attr_create["embeddings"][index]
                    item.update_time = int(time.time())
                await PgAttachmentModel.bulk_create(annex_create)
            if annex_update:
                index: int = 0
                for item in annex_update:
                    item.vector_alias = attr_update["alias"]
                    item.vector_model = attr_update["model"]
                    item.embedding = attr_update["embeddings"][index]
                    item.dimension = len(attr_update["embeddings"][index])
                    item.update_time = int(time.time())
                    await item.save()
            if annex_delete:
                await PgAttachmentModel.filter(uuid__in=annex_delete).update(
                    is_delete=0,
                    delete_time=int(time.time())
                )

            # 更新分词
            await PgKnowledgeModel.update_ts_vectory([str(post.uuid)])
            await PgAttachmentModel.update_ts_vectory(str(post.uuid))

        # 扣费逻辑
        async with in_transaction("mysql"):
            # 用户扣费
            if total_use_points > 0:
                user = await UserModel.filter(id=int(know.id)).first()
                surplus_balance: int = max(user.points - total_use_points, 0)
                await UserModel.filter(id=user.id).update(points=surplus_balance)

            # 写入日志
            if attr_create.get("price") or attr_update.get("price") or ask_emb:
                flows_usage = [flow.model_dump() for flow in additional]
                await UserWalletModel.dec(
                    scene=WalletEnum.POINTS,
                    user_id=int(know.user_id),
                    change_amount=total_use_points,
                    change_type=WalletEnum.POINTS_DEC_KB_IDX,
                    remarks=WalletEnum.get_change_type_msg(WalletEnum.POINTS_DEC_KB_IDX),
                    project=archive.name,
                    source_sn=archive.code,
                    source_id=archive.id,
                    additional=json.dumps(flows_usage, ensure_ascii=False)
                )

    @classmethod
    async def delete(cls, user_id: int, kb: str, uuids: List[str]):
        """
        文档数据删除

        Args:
            user_id (int): 用户ID
            kb (str): 知识库编号
            uuids (List[str]): 数据UUID

        Author:
            zero
        """
        # 验知识库
        know = await AgentKnowModel.filter(code=kb).first()
        if not know:
            raise AppException("知识库已不存在")
        if know.is_delete == 1:
            raise AppException("知识库已被删除")
        if know.is_disable == 1:
            raise AppException("知识库已被禁用")

        # 验证权限
        if int(user_id) != know.user_id:
            await KnowTeamedService.has_permission(know.code, user_id, AgentEnum.POWER_EDIT)

        async with in_transaction("pgsql"):
            # 删除附件
            await PgAttachmentModel.filter(dataset_id__in=uuids).update(
                is_delete=1,
                delete_time=int(time.time())
            )
            # 删除数据
            await (PgKnowledgeModel
                   .filter(know_id=know.id)
                   .filter(uuid__in=uuids, is_delete=0)
                   .update(
                        is_delete=1,
                        last_uid=user_id,
                        delete_time=int(time.time())
                   ))

    @classmethod
    async def status(cls, user_id: int, uuids: str):
        """
        文档数据状态

        Args:
            user_id (int): 用户ID
            uuids (str): 数据UUID

        Author:
            zero
        """
        # 查询数据
        data = await PgKnowledgeModel.filter(uuid=uuids, is_delete=0).get()

        # 验知识库
        know = await AgentKnowModel.filter(id=data.know_id).first()
        if not know:
            raise AppException("知识库已不存在")
        if know.is_delete:
            raise AppException("知识库已被删除")
        if know.is_disable:
            raise AppException("知识库已被禁用")

        # 验证权限
        if int(user_id) != know.user_id:
            await KnowTeamedService.has_permission(know.code, user_id, AgentEnum.POWER_EDIT)

        data.is_disable = 0 if data.is_disable else 1
        await data

    @classmethod
    async def testing(cls, user_id: int, params: schema.KnowDatasetTestingIn) -> List[schema.KnowDatasetTestingVo]:
        """
        数据向量测试

        Args:
            user_id (int): 用户ID
            params (KnowDatasetTestingIn): 检索参数

        Returns:
            List[schema.KnowDatasetTestingVo]

        Author:
            zero
        """
        kb: str = params.kb
        text: str = params.text
        search_mode: str = params.search_mode
        search_tokens: int = params.search_tokens
        search_similar: Decimal = params.search_similar
        ranking_status: bool = bool(params.ranking_status)
        ranking_score: Decimal = params.ranking_score
        ranking_model: str = str(params.ranking_model)
        optimize_ask: bool = bool(params.optimize_ask)
        optimize_model: str = str(params.optimize_model)

        # 验知识库
        know = await AgentKnowModel.filter(code=str(kb)).first()
        if not know:
            raise AppException(str("知识库已不存在"))
        if know.is_disable:
            raise AppException("知识库已被删除")
        if know.is_disable:
            raise AppException("知识库已被禁用")

        # 验证权限
        if user_id != know.user_id:
            await KnowTeamedService.has_permission(know.code, user_id, AgentEnum.POWER_EDIT)

        # 问题优化
        questions = [text]
        if optimize_ask:
            ask = await cls._optimize_ask(text, optimize_model)
            questions.extend(ask)

        # 发起检索
        results = []
        for query in questions:
            if search_mode == "similar":
                lists = await RecallKnowLogic.embedding_recall(
                    model=know.vector_model,
                    query=query,
                    knows=[know.id],
                    limit=100
                )
                results.append({"k": 60, "list": lists})
            elif search_mode == "full":
                lists = await RecallKnowLogic.full_text_recall(
                    query=query,
                    knows=[know.id],
                    limit=100
                )
                results.append({"k": 60, "list": lists})
            elif search_mode == "mix":
                lists = await RecallKnowLogic.mixed_recall(
                    model=know.vector_model,
                    query=query,
                    knows=[know.id]
                )
                results.append({"k": 60, "list": lists})
            else:
                raise AppException("检索模式异常")

        # 结果处理(1): RRF重排
        results = RecallUtil.rrf_concat_results(results)

        # 重排模型
        if ranking_status:
            try:
                data = await RecallUtil.re_ranker_results(results, text, models=ranking_model)
                results = data["results"]
                ranking_status = data["status"]
                if ranking_status:
                    search_similar = ranking_score
            except Exception as e:
                raise AppException(str(e))

        # 相似度过滤
        results = RecallUtil.filter_max_score(
            results=results,
            mode=search_mode,
            rerank=ranking_status,
            similar=search_similar
        )

        # MaxTokens过滤
        results = RecallUtil.filter_max_tokens(results=results, max_tokens=search_tokens)

        # 处理数据
        lists = []
        for item in results:
            lists.append(schema.KnowDatasetTestingVo(
                id=str(item.get("uuid")),
                file=item["file"],
                question=item.get("question", ""),
                answer=item.get("answer", ""),
                character=len(item.get("question", "")) + len(item.get("answer", "")),
                emb_score=item.get("emb_score", -1),
                full_score=item.get("full_score", -1),
                rerank_score=item.get("rerank_score", -1),
            ))

        return lists

    @classmethod
    async def _optimize_ask(cls, question: str, models: str):
        # 查询模型
        try:
            models = await DevModelsModel.check_models(models, scene="chat")
        except Exception as e:
            raise AppException(str(e))

        # 读取模型密钥
        apikey = {}
        if models["config"]["check_key"]:
            apikey = await KeyPoolCache(scene="chat", channel=models["mark"]).get_key() or {}
            if not apikey:
                raise AppException("当前优化模型尚未配置密钥,请与管理员联系")

        # 处理提示词
        template = AgentEnum.get_prompt_tpl("optimize")
        template = template.replace("{{histories}}", "-")
        template = template.replace("{{query}}", question)
        messages = [{"role": "user", "content": template}]

        # 发起对话请求
        aiServer = AIChannelFactory.create(models["channel"], models["config"], apikey)
        results = await aiServer.llm(messages, stream=False)
        content = results.choices[0].message.content

        # 返回问题列表
        return json.loads(content.strip('```json ').strip('```').strip())

    @classmethod
    async def _post_embedding(cls, know: AgentKnowModel, inputs: List[str]):
        # 空的输入
        if not inputs:
            return {
                "price": 0,
                "alias": "",
                "model": "",
                "use_tokens": 0,
                "use_points": 0,
                "embeddings": []
            }

        # 查询模型
        try:
            models = await DevModelsModel.check_models(know.vector_model, ModelsEnum.TYPE_VECTOR)
        except Exception as e:
            raise AppException(str(e))

        # 查询用户
        user = await UserModel.filter(id=know.user_id, is_delete=0).first()
        if not user:
            raise AppException("账号已注销")
        if user.is_disable:
            raise AppException("账号已冻结")
        if user.points <= 0 < models["price"]:
            raise AppException("账号余额不足")

        # 计算费用
        content: str = "".join(inputs)
        total_tokens: int = ChatUtils.compute_tokens(content)
        total_points: int = ChatUtils.compute_price(total_tokens, models["price"])
        if models["price"] and user.points < total_points:
            raise AppException("账户余额不足")

        # 内容向量化
        vectorService = VectorService()
        embedding_arr = await vectorService.to_embed(models, inputs)

        # 返回结果
        return {
            "price": models["price"],
            "alias": models["alias"],
            "model": models["model"],
            "use_tokens": total_tokens,
            "use_points": total_points,
            "embeddings": embedding_arr
        }
